import tensorflow as tf
import numpy as np

class Client:
  def __init__(self, model_factory, data, label, learning_rate, R, batch_size):
    self.attacker = False               # 是否是攻击者，默认为 False
    self.threat_model = None            # 攻击模型（如果客户端是攻击者，则会设置该值）
    self._x = data                      # 存储本地数据（特征）
    self._y = label                     # 存储本地数据（标签）
    self._model = model_factory()       # 通过 model_factory() 创建客户端本地模型
    self.learning_rate = learning_rate  # 记录学习率
    self.epochs = R                     # 本地训练轮数
    self.batch_size = batch_size        # 本地训练时的小批量大小
    self.steps_per_epoch = 1            # 训练步数（默认为 1，可根据数据大小调整）
    # 重新编译模型，使用 SGD（随机梯度下降）优化器，交叉熵损失函数，并评估准确率
    self._model.compile(
      optimizer=tf.keras.optimizers.SGD(learning_rate=self.learning_rate),
      loss=tf.keras.losses.CategoricalCrossentropy(),
      metrics=['accuracy']
    )

  # 接收服务器下发的全局权重
  def train(self, server_weights):
    # 使用服务器权重初始化本地模型
    self._model.set_weights(server_weights)
    # 在本地数据上执行训练（epochs=R，默认 R=1）
    self._model.fit(x = self._x, y = self._y, verbose=0,
                    epochs = self.epochs, batch_size = self.batch_size,
                    steps_per_epoch=self.steps_per_epoch,
                    )
    # 获取本地训练后的新模型权重
    new_weights = self._model.get_weights()
    # 计算权重的变化量（delta_weights = 新权重 - 服务器权重）
    delta_weights = [new_w - old_w for new_w, old_w in zip(new_weights, server_weights)]
    # 如果该客户端是攻击者，可能需要篡改 delta_weights
    return delta_weights